{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "da5198c6",
   "metadata": {},
   "source": [
    "(sec-dataset-advanced-examples)=\n",
    "# Quantify dataset - advanced examples\n",
    "\n",
    "```{seealso}\n",
    "The complete source code of this tutorial can be found in\n",
    "\n",
    "{nb-download}`Quantify dataset - advanced examples.ipynb`\n",
    "```\n",
    "\n",
    "Here we will explore a few advanced usages of the Quantify dataset and how it can\n",
    "accommodate them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64643c33",
   "metadata": {
    "mystnb": {
     "code_prompt_show": "Imports and auxiliary utilities"
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "from rich import pretty\n",
    "\n",
    "from quantify_core.analysis.calibration import rotate_to_calibrated_axis\n",
    "from quantify_core.analysis.fitting_models import exp_decay_func\n",
    "from quantify_core.data import handling as dh\n",
    "from quantify_core.utilities import dataset_examples\n",
    "from quantify_core.utilities.dataset_examples import (\n",
    "    mk_nested_mc_dataset,\n",
    "    mk_shots_from_probabilities,\n",
    "    mk_surface7_cyles_dataset,\n",
    ")\n",
    "from quantify_core.utilities.examples_support import (\n",
    "    mk_iq_shots,\n",
    "    round_trip_dataset,\n",
    ")\n",
    "from quantify_core.utilities.inspect_utils import display_source_code\n",
    "\n",
    "pretty.install()\n",
    "\n",
    "dh.set_datadir(Path.home() / \"quantify-data\")  # change me!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b075ee74",
   "metadata": {},
   "source": [
    "## Dataset for an \"unstructured\" experiment\n",
    "\n",
    "Let's take consider a Surface Code experiment, in particular the one portrayed in\n",
    "Fig. 4b from one of the papers from DiCarlo Lab {cite}`marques_logical_qubit_2021`.\n",
    "\n",
    "For simplicity, we will not use exactly the same schedule, because what matters here\n",
    "are the measurements. It is difficult to deal with the results of these measurements\n",
    "because we have a few repeating cycles followed by a final measurement that leaves the\n",
    "overall dataset \"unstructured\".\n",
    "\n",
    "```{figure} /images/surface-7-sched.png\n",
    ":width: 100%\n",
    "```\n",
    "\n",
    "``````{admonition} Source code for generating this schedule and visualizing it\n",
    ":class: dropdown, info\n",
    "\n",
    "If you want to create and visualize the schedule above using `quantify-scheduler`,\n",
    "you can use this code:\n",
    "\n",
    "```{code-block} python\n",
    "from quantify_scheduler import Schedule\n",
    "from quantify_scheduler.operations.gate_library import CZ, Y90, Measure, Reset, X\n",
    "from quantify_scheduler.visualization.circuit_diagram import circuit_diagram_matplotlib\n",
    "\n",
    "def mk_surface7_sched(num_cycles: int = 3):\n",
    "    \"\"\"Generates a schedule with some of the feature of a Surface 7 experiment as\n",
    "    portrayed in Fig. 4b of :cite:`marques_logical_qubit_2021`.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    num_cycles\n",
    "        The number of times to repeat the main cycle.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    :\n",
    "        A schedule similar to a Surface 7 dance.\n",
    "    \"\"\"\n",
    "\n",
    "    sched = Schedule(\"S7 dance\")\n",
    "\n",
    "    q_d1, q_d2, q_d3, q_d4 = [f\"D{i}\" for i in range(1, 5)]\n",
    "    q_a1, q_a2, q_a3 = [f\"A{i}\" for i in range(1, 4)]\n",
    "    all_qubits = q_d1, q_d2, q_d3, q_d4, q_a1, q_a2, q_a3\n",
    "\n",
    "    sched.add(Reset(*all_qubits))\n",
    "\n",
    "    for cycle in range(num_cycles):\n",
    "        sched.add(Y90(q_d1))\n",
    "        for qubit in [q_d2, q_d3, q_d4]:\n",
    "            sched.add(Y90(qubit), ref_pt=\"start\", rel_time=0)\n",
    "        sched.add(Y90(q_a2), ref_pt=\"start\", rel_time=0)\n",
    "\n",
    "        for qubit in [q_d2, q_d1, q_d4, q_d3]:\n",
    "            sched.add(CZ(qC=qubit, qT=q_a2))\n",
    "\n",
    "        sched.add(Y90(q_d1))\n",
    "        for qubit in [q_d2, q_d3, q_d4]:\n",
    "            sched.add(Y90(qubit), ref_pt=\"start\", rel_time=0)\n",
    "        sched.add(Y90(q_a2), ref_pt=\"start\", rel_time=0)\n",
    "\n",
    "        sched.add(Y90(q_a1), ref_pt=\"end\", rel_time=0)\n",
    "        sched.add(Y90(q_a3), ref_pt=\"start\", rel_time=0)\n",
    "\n",
    "        sched.add(CZ(qC=q_d1, qT=q_a1))\n",
    "        sched.add(CZ(qC=q_d2, qT=q_a3))\n",
    "        sched.add(CZ(qC=q_d3, qT=q_a1))\n",
    "        sched.add(CZ(qC=q_d4, qT=q_a3))\n",
    "\n",
    "        sched.add(Y90(q_a1), ref_pt=\"end\", rel_time=0)\n",
    "        sched.add(Y90(q_a3), ref_pt=\"start\", rel_time=0)\n",
    "\n",
    "        sched.add(Measure(q_a2, acq_index=cycle))\n",
    "        for qubit in (q_a1, q_a3):\n",
    "            sched.add(Measure(qubit, acq_index=cycle), ref_pt=\"start\", rel_time=0)\n",
    "\n",
    "        for qubit in [q_d1, q_d2, q_d3, q_d4]:\n",
    "            sched.add(X(qubit), ref_pt=\"start\", rel_time=0)\n",
    "\n",
    "    # final measurements\n",
    "\n",
    "    sched.add(Measure(*all_qubits[:4], acq_index=0), ref_pt=\"end\", rel_time=0)\n",
    "\n",
    "    return sched\n",
    "\n",
    "sched = mk_surface7_sched(num_cycles=3)\n",
    "f, ax = circuit_diagram_matplotlib(sched)\n",
    "f.set_figwidth(30)\n",
    "```\n",
    "``````\n",
    "\n",
    "How do we store all the shots for this measurement?\n",
    "We might want this because, e.g., we know we have an issue with leakage to the second\n",
    "excited state of a transmon and we would like to be able to store and inspect raw data.\n",
    "\n",
    "To support such use-cases we will have a dimension in the dataset for the repeating cycles\n",
    "and one extra dimension for the final measurement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2f2ddcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# mock data parameters\n",
    "num_shots = 128  # NB usually >~1000 in real experiments\n",
    "ground = -0.2 + 0.65j\n",
    "excited = 0.7 + 4j\n",
    "centroids = ground, excited\n",
    "sigmas = [0.1] * 2\n",
    "\n",
    "display_source_code(mk_iq_shots)\n",
    "display_source_code(mk_shots_from_probabilities)\n",
    "display_source_code(mk_surface7_cyles_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12829358",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = mk_surface7_cyles_dataset(\n",
    "    num_shots=num_shots, sigmas=sigmas, centers=centroids\n",
    ")\n",
    "\n",
    "assert dataset == round_trip_dataset(dataset)  # confirm read/write\n",
    "\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b71d9599",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.A1_shots.shape, dataset.D1_shots.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cf70976",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_gridded = dh.to_gridded_dataset(\n",
    "    dataset, dimension=\"dim_cycle\", coords_names=[\"cycle\"]\n",
    ")\n",
    "dataset_gridded = dh.to_gridded_dataset(\n",
    "    dataset_gridded, dimension=\"dim_final\", coords_names=[\"final_msmt\"]\n",
    ")\n",
    "dataset_gridded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a216dc93",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_gridded.A0_shots.real.mean(\"repetitions\").plot(marker=\"o\", label=\"I-quadrature\")\n",
    "dataset_gridded.A0_shots.imag.mean(\"repetitions\").plot(marker=\"^\", label=\"Q-quadrature\")\n",
    "_ = plt.gca().legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41b4f6af",
   "metadata": {},
   "source": [
    "(sec-nested-mc-example)=\n",
    "## Dataset for a \"nested MeasurementControl\" experiment\n",
    "\n",
    "Now consider a dataset that has been constructed by an experiment involving the\n",
    "operation of two\n",
    "{class}`.MeasurementControl` objects. The second of\n",
    "them performs a \"meta\" outer loop in which we sweep a flux bias and then perform\n",
    "several experiments to characterize a transmon qubit, e.g. determining the frequency of\n",
    "a read-out resonator, the frequency of the transmon, and its T1 lifetime.\n",
    "\n",
    "Below we showcase what the data from the dataset containing the T1 experiment results\n",
    "could look like"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d839806c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "rng = np.random.default_rng(seed=112244)  # random number generator\n",
    "\n",
    "num_t1_datasets = 7\n",
    "t1_times = np.linspace(0, 120e-6, 30)\n",
    "\n",
    "for tau in rng.uniform(10e-6, 50e-6, num_t1_datasets):\n",
    "    probabilities = exp_decay_func(\n",
    "        t=t1_times, tau=tau, offset=0, n_factor=1, amplitude=1\n",
    "    )\n",
    "    dataset = dataset_examples.mk_t1_av_with_cal_dataset(t1_times, probabilities)\n",
    "\n",
    "    round_trip_dataset(dataset)  # confirm read/write\n",
    "    dataset_g = dh.to_gridded_dataset(\n",
    "        dataset, dimension=\"main_dim\", coords_names=[\"t1_time\"]\n",
    "    )\n",
    "    # rotate the iq data\n",
    "    rotated_and_normalized = rotate_to_calibrated_axis(\n",
    "        dataset_g.q0_iq_av.values, *dataset_g.q0_iq_av_cal.values\n",
    "    )\n",
    "    rotated_and_normalized_da = xr.DataArray(dataset_g.q0_iq_av)\n",
    "    rotated_and_normalized_da.values = rotated_and_normalized\n",
    "    rotated_and_normalized_da.attrs[\"long_name\"] = \"|1> Population\"\n",
    "    rotated_and_normalized_da.attrs[\"units\"] = \"\"\n",
    "    rotated_and_normalized_da.real.plot(ax=ax, label=dataset.tuid, marker=\".\")\n",
    "ax.set_title(\"Results from repeated T1 experiments\\n(different datasets)\")\n",
    "_ = ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae491ce5",
   "metadata": {},
   "source": [
    "Since the raw data is now split among several datasets, we would like to keep a\n",
    "reference to all these datasets in our \"combined\" datasets. Below we showcase how this\n",
    "can be achieved, along with some useful xarray features and known limitations.\n",
    "\n",
    "We start by generating a mock dataset that combines all the information that would have\n",
    "been obtained from analyzing a series of other datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7149f12",
   "metadata": {
    "mystnb": {
     "code_prompt_show": "Source code for mk_nested_mc_dataset function"
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "display_source_code(mk_nested_mc_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09512820",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = mk_nested_mc_dataset(num_points=num_t1_datasets)\n",
    "assert dataset == round_trip_dataset(dataset)  # confirm read/write\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05758f71",
   "metadata": {},
   "source": [
    "In this case the four main coordinates are not orthogonal coordinates, but instead\n",
    "just different labels for the same data points, also known as a \"multi-index\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8a4264e",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(3, 1, figsize=(10, 10), sharex=True)\n",
    "\n",
    "_ = dataset.t1.plot(x=\"flux_bias\", marker=\"o\", ax=axs[0].twiny(), color=\"C0\")\n",
    "x = \"t1_tuids\"\n",
    "_ = dataset.t1.plot(x=x, marker=\"o\", ax=axs[0], color=\"C0\")\n",
    "_ = dataset.resonator_freq.plot(x=x, marker=\"o\", ax=axs[1], color=\"C1\")\n",
    "_ = dataset.qubit_freq.plot(x=x, marker=\"o\", ax=axs[2], color=\"C2\")\n",
    "for tick in axs[2].get_xticklabels():\n",
    "    tick.set_rotation(15)  # avoid tuid labels overlapping"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92c30a1b",
   "metadata": {},
   "source": [
    "It is possible to work with an explicit MultiIndex within a (python) xarray object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d1df2bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_multi_indexed = dataset.set_index({\"main_dim\": tuple(dataset.t1.coords.keys())})\n",
    "dataset_multi_indexed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15bbd715",
   "metadata": {},
   "source": [
    "The MultiIndex is very handy for selecting data in different ways, e.g.:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dc455da",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = 2\n",
    "dataset_multi_indexed.qubit_freq.sel(\n",
    "    qubit_freq_tuids=dataset_multi_indexed.qubit_freq_tuids.values[index]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37a714fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_multi_indexed.qubit_freq.sel(t1_tuids=dataset.t1_tuids.values[index])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6fd821c",
   "metadata": {},
   "source": [
    "### Known limitations\n",
    "\n",
    "Unfortunately, at the moment the MultiIndex has the problem of not being compatible with\n",
    "the NetCDF format used to write to disk:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cdfdf30",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    assert dataset_multi_indexed == round_trip_dataset(dataset_multi_indexed)\n",
    "except NotImplementedError as exp:\n",
    "    print(exp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2596e9ca",
   "metadata": {},
   "source": [
    "We could make our load/write utilities take care of setting and resetting the index\n",
    "under the hood. Though there are some nuances there as well. If we would do that then\n",
    "some extra metadata needs to be stored in order to store/restore the multi-index.\n",
    "At the moment, the MultiIndex is not supported when writing a Quantify dataset to\n",
    "disk. Below we show a few complications related to this.\n",
    "\n",
    "Fortunately, the MultiIndex can be reset back:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcdb4368",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_multi_indexed.reset_index(dims_or_levels=\"main_dim\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "985b7c82",
   "metadata": {},
   "outputs": [],
   "source": [
    "all(dataset_multi_indexed.reset_index(\"main_dim\").t1_tuids == dataset.t1_tuids)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27741c5c",
   "metadata": {},
   "source": [
    "But, for example, the `dtype` has been changed to `object`\n",
    "(from fixed-length string):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62cb3085",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.t1_tuids.dtype, dataset_multi_indexed.reset_index(\"main_dim\").t1_tuids.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db2b5eaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.t1_tuids.dtype == dataset_multi_indexed.reset_index(\"main_dim\").t1_tuids.dtype"
   ]
  }
 ],
 "metadata": {
  "file_format": "mystnb",
  "jupytext": {
   "text_representation": {
    "extension": ".md",
    "format_name": "myst"
   }
  },
  "kernelspec": {
   "display_name": "python3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}